{
 "cells": [
  {
   "cell_type": "code",
   "execution_count": 1,
   "metadata": {},
   "outputs": [
    {
     "data": {
      "text/plain": [
       "'/AIRvePFS/dair/luoyz-data/Open_DAIR_BioMed'"
      ]
     },
     "execution_count": 1,
     "metadata": {},
     "output_type": "execute_result"
    }
   ],
   "source": [
    "import os.path as osp\n",
    "import sys\n",
    "path = osp.dirname(osp.abspath(''))\n",
    "sys.path.append(path)\n",
    "sys.path.append(osp.join(path, \"open_biomed\"))\n",
    "path"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 2,
   "metadata": {},
   "outputs": [
    {
     "name": "stderr",
     "output_type": "stream",
     "text": [
      "/AIRvePFS/dair/luoyz-data/anaconda3/envs/biomedgpt/lib/python3.9/site-packages/tqdm/auto.py:21: TqdmWarning: IProgress not found. Please update jupyter and ipywidgets. See https://ipywidgets.readthedocs.io/en/stable/user_install.html\n",
      "  from .autonotebook import tqdm as notebook_tqdm\n"
     ]
    },
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Config:  {'name': 'biomedgptv', 'data': {'mol': {'modality': ['structure'], 'featurizer': {'structure': {'name': 'BaseGNN'}}}, 'protein': {'modality': ['structure'], 'featurizer': {'structure': {'name': 'transformertok', 'transformer_type': 'esm', 'model_name_or_path': '/AIRvePFS/dair/luoyz-data/Open_DAIR_BioMed/ckpts/protein_ckpts/esm2-3b', 'max_length': 1024}}}}, 'network': {'mol': {'gin_hidden_dim': 300, 'gin_num_layers': 5, 'drop_ratio': 0.0, 'max_n_nodes': 256, 'freeze': True}, 'protein': {'model_name_or_path': '/AIRvePFS/dair/luoyz-data/Open_DAIR_BioMed/ckpts/protein_ckpts/esm2-3b', 'lora': False, 'freeze': True, 'use_float16': True}, 'llm': {'model_name_or_path': '/AIRvePFS/dair/luoyz-data/Open_DAIR_BioMed/ckpts/text_ckpts/biomed-llama-7b', 'use_float16': True}}}\n",
      "Finish loading model\n"
     ]
    }
   ],
   "source": [
    "# load biomedgpt model\n",
    "import json\n",
    "import torch\n",
    "from open_biomed.utils import fix_path_in_config\n",
    "from open_biomed.models.multimodal import BioMedGPTV\n",
    "\n",
    "config = json.load(open(\"../configs/encoders/multimodal/biomedgptv.json\", \"r\"))\n",
    "fix_path_in_config(config, path)\n",
    "print(\"Config: \", config)\n",
    "\n",
    "# ~20G GPU Memory is required to load BioMedGPT-10B\n",
    "device = torch.device(\"cuda:1\")\n",
    "config[\"network\"][\"device\"] = device\n",
    "model = BioMedGPTV(config[\"network\"])\n",
    "ckpt = torch.load(\"../ckpts/fusion_ckpts/biomedgpt_10b.pth\")\n",
    "model.load_state_dict(ckpt)\n",
    "model = model.to(device)\n",
    "model.eval()\n",
    "print(\"Finish loading model\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 5,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Human:  Please describe this molecule.\n",
      "Assistant:  The molecule is a member of the class of flavonoids that is 3-hydroxyflavone substituted by a 3-hydroxy-4-methoxyphenyl group at position 7. It has a role as a plant metabolite. It is a member of flavonoids, a member of 3-hydroxyflavones and a member of 3-hydroxyflavones. It derives from a 3-hydroxyflavone.\n"
     ]
    }
   ],
   "source": [
    "# chat\n",
    "from open_biomed.utils.chat_utils import Conversation\n",
    "\n",
    "prompt_sys = \"You are working as an excellent assistant in chemistry and molecule discovery. \" + \\\n",
    "             \"Below a human expert gives the representation of a molecule or a protein. Answer questions about it. \"\n",
    "chat = Conversation(\n",
    "    model=model, \n",
    "    processor_config=config[\"data\"], \n",
    "    device=device,\n",
    "    system=prompt_sys,\n",
    "    roles=(\"Human\", \"Assistant\"),\n",
    "    sep=\"###\",\n",
    "    max_length=2048\n",
    ")\n",
    "chat.append_molecule(\"CC(=CCC1=CC(=CC(=C1O)CC=C(C)C)/C=C/C(=O)C2=C(C=C(C=C2)O)O)C\")\n",
    "question = \"Please describe this molecule.\"\n",
    "print(\"Human: \", question)\n",
    "chat.ask(question)\n",
    "print(\"Assistant: \", chat.answer()[0])"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 6,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Human:  What is the function of this protein?\n",
      "Assistant:  One of the essential components for the initiation of protein synthesis.  Stabilizes the binding of IF-2 and IF-3 on the 30S subunit to which N-formylmethionyl-tRNA(fMet) subsequently binds.  Helps modulate mRNA selection, yielding the 30S pre-initiation complex (PIC).  Upon addition of the 50S ribosomal subunit IF-1, IF-2 and IF-3 are released leaving the mature 70S translation initiation complex.\n"
     ]
    }
   ],
   "source": [
    "chat.reset()\n",
    "chat.append_protein(\"MAKEDTLEFPGVVKELLPNATFRVELDNGHELIAVMAGKMRKNRIRVLAGDKVQVEMTPYDLSKGRINYRFK\")\n",
    "question = \"What is the function of this protein?\"\n",
    "print(\"Human: \", question)\n",
    "chat.ask(question)\n",
    "print(\"Assistant: \", chat.answer()[0])"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "Python 3.9.17 ('biomedgpt')",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.9.17"
  },
  "orig_nbformat": 4,
  "vscode": {
   "interpreter": {
    "hash": "c9acc99eda29845eb06fd58aaf0629edb79f58e739da13f56214d4672f20527d"
   }
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
